from PIL import Image
import sys
from PIL import ImageGrab
import keyboard
import time
from aip import AipOcr
import json
import os

# 副标题的间距 type == 1
from computedComposing import computedCol, computedRow, setGlobalValue, computedComposing

SUB_SPACE = 20
# 同一行top值的相差精度 , 用于判定是否在同一行
TOP_PRECISION = 10
# 同一行left值的差精度 , 用于判定是否是同一段文字 , 被空格分开
LEFT_PRECISION = 10
"""
模式设置
0 上下表单模式
1 横向左右数据模式
2 普通模式
"""
TYPE = 0
# 截屏热键
SCREEN_HOTKEY = 'ctrl+shift+a'
# 截图后 , 开始处理功能热键
START_HOTKEY = 'ctrl+1'


# 获取控制台的参数
def getArgs():
    i = 0
    avg = sys.argv[1:]
    while i < len(avg):
        key = avg[i].lower()
        val = avg[i + 1]
        print(key, val)
        if key == '-ss' or key == '-sub_space':
            global SUB_SPACE
            SUB_SPACE = int(val)
        elif key == '-tp' or key == '-TOP_PRECISION'.lower():
            global TOP_PRECISION
            TOP_PRECISION = int(val)
        elif key == '-lp' or key == '-LEFT_PRECISION'.lower():
            global LEFT_PRECISION
            LEFT_PRECISION = int(val)
        elif key == '-t' or key == '-TYPE'.lower():
            global TYPE
            TYPE = int(val)
        elif key == '-s' or key == '-SCREEN_HOTKEY'.lower():
            global SCREEN_HOTKEY
            SCREEN_HOTKEY = val
        elif key == '-sh' or key == '-START_HOTKEY'.lower():
            global START_HOTKEY
            START_HOTKEY = val
        i += 2
    setGlobalValue(TOP_PRECISION,LEFT_PRECISION)

getArgs()


def screenshot(client):
    keyboard.wait(SCREEN_HOTKEY)
    print("截图")
    keyboard.wait(START_HOTKEY)  # 输入键盘的触发事件
    print("截图完成")
    # 2.保存图片
    image = ImageGrab.grabclipboard()  # 把图片从剪切板保存到当前路径
    image.save("screen.png")
    time.sleep(0.1)
    handle_baidu_aip(client, TYPE)


# 处理百度api
def handle_baidu_aip(client, type):
    with open("screen.png", 'rb') as f:
        image = f.read()
        # 调用百度API通用文字识别（高精度版），提取图片中的内容
        # text = client.basicAccurate(image)
        # 调用通用文字识别（高精度含位置版）
        text = client.accurate(image)
        result = text["words_result"]
        count = text['words_result_num']
        file = open(os.getcwd() + '\\result.json', 'w+', encoding='utf-8')
        file.write(json.dumps(result, indent=4, ensure_ascii=False))
        file.close()
        res = []
        if type == 0:
            res = handle_type_0_ex(result, count)
        elif type == 1:
            res = handle_type_1(result, count)
        elif type == 2:
            res = handle_type_2(result, count)
        toJson = json.dumps(res, indent=4, ensure_ascii=False)
        file = open(os.getcwd() + '\\data.json', 'w+', encoding='utf-8')
        file.write(toJson)
        file.close()
        print("写入成功!")


# 处理type 0
def handle_type_0(result, count):
    myJson = []
    jumpNextLoop = False
    for i in range(count):
        if jumpNextLoop:
            jumpNextLoop = False
            continue
        item = result[i]
        type = ""
        nex = False if i + 1 >= count else result[i + 1]['words']
        isTitle = True
        ph = ""
        if '>' in item['words'] or '7' == item['words']:
            continue
        if nex:
            if '请输入' in nex:
                type = 'input'
                isTitle = False
                jumpNextLoop = True
                ph = nex
            elif '请选择' in nex:
                type = 'select'
                isTitle = False
                ph = nex
                jumpNextLoop = True

        obj = {
            "key": item["words"],
            "type": type,
            "placeholder": ph,
            "local": item['location']
        }
        if type == 'select': obj['options'] = []
        myJson.append(obj)
    return handleSubtitle(myJson)

def handle_type_0_ex(result, count):
    myJson = []
    i = 0
    while i < count:
        # 本条记录
        item = result[i]
        # 下一条记录
        nextOne = False if i + 1 >= count else result[i + 1]
        # 获取排版信息
        composing = computedComposing(item,nextOne)
        # 判定下一个字段是否是和当前字段是同一个字段
        if nextOne and composing[1] < LEFT_PRECISION and composing[0] <= 5:
            item['words'] += nextOne['words']
            del result[i+1]
            count -= 1
            nextOne = False if i + 1 >= count else result[i + 1]
            composing = computedComposing(item, nextOne)
        addI = 1
        res = {}
        # 计算类型
        if composing[2]:
            res = computedCol(result,i,composing[0])
        else:
            res = computedRow(result,i,composing[1])
        addI += res['jmp']
        i = i + addI
        del res['jmp']
        myJson.append(res)
    return myJson

# 处理 type 1
def handle_type_1(result, count):
    jumpNextLoop = False
    temp = []
    for i in range(count):
        item = result[i]
        if i == 0:
            temp.append({
                "key": item['words'],
                'isTitle': False,
                'location': item['location']
            })
            continue
        prev = temp.pop()
        cur = {
            "key": item['words'],
            'isTitle': False,
            'location': item['location']
        }
        # top数据
        pt = prev['location']['top']
        ct = item['location']['top']

        pl = prev['location']['left'] + prev['location']['width']  # 上一条的right数据
        cl = item['location']['left']  # 当前的left数据
        # 如果不在同一行
        if abs(pt - ct) > TOP_PRECISION:
            if not prev.get('val'):
                prev['isTitle'] = True
            else:
                prev['isTitle'] = False
            temp.append(prev)
            temp.append(cur)
        # 如果在同一行 , 且不是同一组数据
        elif abs(pl - cl) > LEFT_PRECISION:
            if not prev.get('val'):
                prev['val'] = [item['words']]
            else:
                prev['val'].append(item['words'])
            temp.append(prev)
    return temp


# 处理 type 2
def handle_type_2(result, count):
    res = []
    for item in result:
        res.append(item['words'])
    return res


# 处理副标题
def handleSubtitle(obj):
    temp = []
    for i in range(len(obj)):
        item = obj[i]
        if i == len(obj) - 1:
            temp.append(item)
            break
        ne = obj[i + 1]
        curTop = item['local']['top']
        curH = item['local']['height']
        nextTop = ne['local']['top']
        count = nextTop - curTop - curH
        if SUB_SPACE >= count:
            ne['isST'] = True
        temp.append(item)
    i = 0
    res = []
    while i < len(temp):
        item = temp[i]
        item.pop("local")
        if not not item.get('isST'):
            last = res.pop()
            if not not last.get('subTitle'):
                last['subTitle'].append(item['key'])
            else:
                last['subTitle'] = [item['key']]
            res.append(last)
        else:
            res.append(item)
        i += 1
    return res


if __name__ == '__main__':
    APP_ID = '27749896'
    API_KEY = 'q800E896WenTUeGT9oRj0Iv3'
    SECRET_KEY = '5Tfrzc8cqujdbnCqYrOirralkUza5Dm2'

    client = AipOcr(APP_ID, API_KEY, SECRET_KEY)
    while True:
        screenshot(client)
